# 运行前请 pip install tiktoken, python-docx
from dashscope import get_tokenizer  # dashscope版本 >= 1.14.0
import docx
import os

def read_txt_file(file_path):
    """读取txt文件内容"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return f.read()

def read_docx_file(file_path):
    """读取docx文件内容"""
    doc = docx.Document(file_path)
    return '\n'.join([para.text for para in doc.paragraphs])

def count_tokens(text, tokenizer):
    """统计文本的token数量"""
    tokens = tokenizer.encode(text)
    return len(tokens), tokens

def main():
    # 获取tokenizer对象，目前只支持通义千问系列模型
    tokenizer = get_tokenizer('qwen-plus')
    
    # 示例字符串
    input_str = '通义千问具有强大的能力。'
    token_count, tokens = count_tokens(input_str, tokenizer)
    print(f"示例字符串: '{input_str}'")
    print(f"经过切分后的token id为：{tokens}。")
    print(f"经过切分后共有{token_count}个token\n")
    
    # 读取文档文件
    file_path = r"D:\Project\Test\utils\中华人民共和国立法法.docx"
    
    if not os.path.exists(file_path):
        print("文件不存在!")
        return
    
    if file_path.endswith('.txt'):
        content = read_txt_file(file_path)
    elif file_path.endswith('.docx'):
        content = read_docx_file(file_path)
    else:
        print("不支持的文件格式，请提供.txt或.docx文件")
        return
    
    # 统计token
    token_count, tokens = count_tokens(content, tokenizer)
    
    print("\n文档内容分析结果:\n================================================")
    print(f"文档内容长度(字符): {len(content)}")
    print(f"文档内容token数量: {token_count}")
    
    # 可选：显示部分token详情
    show_details = input("\n是否要显示前20个token的详情?(y/n): ").lower()
    if show_details == 'y':
        print("\n前20个token详情:")
        for i in range(min(20, len(tokens))):
            print(f"token id为{tokens[i]}对应的字符串为：{tokenizer.decode(tokens[i])}")

if __name__ == '__main__':
    main()